import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test setting a distribution of models
for several agents, creating different possible worlds
'''

class MultipleEvilTwinsTestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.maxRounds=5
        self.world = World()
        self.trevor = Agent('Trevor')
        self.world.addAgent(self.trevor)
        self.camille = Agent('Camille')
        self.world.addAgent(self.camille)
        
        self.agents = [self.trevor, self.camille]
        
        self.world.setOrder(['Trevor', 'Camille'])        

    def addStates(self):
        #define states        
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        
        for agent in self.agents:
            self.world.defineState(agent.name,'karma',int)
            self.world.setState(agent.name,'karma',10)
        
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))

    def addActions(self):
        #add actions and appropriate dynamics
        for agent in self.agents:           
            key = stateKey(agent.name, 'karma')
            action =  agent.addAction({'verb': 'do good'})
            tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
            self.world.setDynamics('round',action,tree) 
            tree = makeTree(incrementMatrix(key,1))
            self.world.setDynamics(key,action,tree) 
        
            action =  agent.addAction({'verb': 'do evil'})
            tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
            self.world.setDynamics('round',action,tree) 
            tree = makeTree(incrementMatrix(key,-1))
            self.world.setDynamics(key,action,tree)
               
    def setModelAndRewards(self):
        #add good/evil models
        self.trevor.addModel('GoodTrevor',R={},rationality=1.,level=2)
        self.trevor.addModel('EvilTrevor',R={},rationality=1.,level=2)
        
        self.trevor.setReward(maximizeFeature(stateKey('Trevor', 'karma')), 1., 'GoodTrevor')
        self.trevor.setReward(minimizeFeature(stateKey('Trevor', 'karma')), 1., 'EvilTrevor')
        
        beliefs = {'GoodTrevor':0.7, 'EvilTrevor':0.3}
         
        self.world.setModel('Trevor', beliefs)

        self.camille.addModel('GoodCamille',R={},rationality=1.,level=2)
        self.camille.addModel('EvilCamille',R={},rationality=1.,level=2)
        
        self.camille.setReward(maximizeFeature(stateKey('Camille', 'karma')), 1., 'GoodCamille')
        self.camille.setReward(minimizeFeature(stateKey('Camille', 'karma')), 1., 'EvilCamille')
        
        beliefs = {'GoodCamille':0.4, 'EvilCamille':0.6}
         
        self.world.setModel('Camille', beliefs)                                                          

    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.setModelAndRewards()
        
    def testRun(self):
        self.testSetup()
        self.world.printState()
        self.world.printModels()
        for t in range(self.maxRounds + 1):
            self.world.explain(self.world.step(),level=1)
#             self.world.state.select()
            self.world.printState()
            self.world.printModels()
            if self.world.terminated():
                break 
            
if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(MultipleEvilTwinsTestCase):
       if method.startswith("test"):
          suite.addTest(MultipleEvilTwinsTestCase(method))
    unittest.TextTestRunner().run(suite)